//
// memb_listener.h
//

#ifndef KITTY_MEMB_LISTENER_H
#define KITTY_MEMB_LISTENER_H

#include <boost/asio.hpp>
#include <boost/bind.hpp>
#include <boost/thread.hpp>
#include <boost/tuple/tuple.hpp>
#include <boost/lexical_cast.hpp>

#include "util/MD5.h"
#include "heartbeat.h"
#include "soft_state.h"
#include "mcast_receiver.h"

namespace kitty {

class memb_listener
{
public:
    memb_listener(boost::asio::io_service& io_service,
        boost::shared_mutex& soft_state_mutex,
        kitty::soft_state& soft_state,
        const boost::asio::ip::address& listen_address,
        const boost::asio::ip::address& mcast_address,
        unsigned short mcast_port)
      : soft_state_mutex_(soft_state_mutex),
        soft_state_(soft_state),
        mcast_receiver_(io_service, listen_address, mcast_address, mcast_port)
    {
        kitty::heartbeat_ptr h_ptr(new kitty::heartbeat());
        mcast_receiver_.async_read(
                *h_ptr, boost::bind(
                    &memb_listener::handle_receive_hearbeat, this,
                    boost::asio::placeholders::error, h_ptr));
    }

private:
    void handle_receive_hearbeat(
            const boost::system::error_code& e, kitty::heartbeat_ptr heartbeat)
    {
        if (!e)
        {
            // timestamp when it receives this heartbeat packet
            time_t local_time = time(NULL);
            // heartbeat is copied into the internal of the  new thread
            boost::thread t(
                    boost::bind(&memb_listener::process_heartbeat, this, 
                        heartbeat, local_time));

            kitty::heartbeat_ptr h_ptr(new kitty::heartbeat());
            mcast_receiver_.async_read(
                    *h_ptr, boost::bind(
                        &memb_listener::handle_receive_hearbeat, this,
                        boost::asio::placeholders::error, h_ptr));
        }
        else
        {
            std::cerr << "Something went wrong in mcast_receiver: ";
            std::cerr << e.message() << std::endl;
            exit(1);
        }
    }

    void process_heartbeat(kitty::heartbeat_ptr heartbeat, time_t local_time)
    {
        // MD5(ip:port) is the key of the soft state map
        std::ostringstream vnode_stream;
        vnode_stream << heartbeat->ip << ":" << heartbeat->port;
        std::string vnode = vnode_stream.str();
        kitty::util::MD5 md5(
                reinterpret_cast<const unsigned char*>(vnode.c_str()),
                static_cast<int>(vnode.size()));
        boost::array<unsigned char, 16> key;
        md5.digest(key.c_array());

        boost::shared_lock<boost::shared_mutex> rlock(soft_state_mutex_);

        if (soft_state_.find(key) == soft_state_.end())
        {
            rlock.unlock();
            boost::unique_lock<boost::shared_mutex> wlock(soft_state_mutex_);

            /* Problem
             * What if the writer is starved? The priority of the writer
             * should be higher than the reader. */
            if (soft_state_.find(key) == soft_state_.end())
            {
                kitty::soft_state_entry_ptr entry(new kitty::soft_state_entry());

                entry->local_time = local_time;
                entry->ip = heartbeat->ip;
                entry->port = heartbeat->port;
                entry->peer_time = heartbeat->timestamp;
                entry->state = kitty::peer_alive;

                soft_state_.insert(soft_state::value_type(key, entry));
            }
        }
        else
        {
            boost::lock_guard<boost::mutex> guard(soft_state_[key]->entry_mutex);
            if (soft_state_[key]->state == kitty::peer_alive
                    && soft_state_[key]->peer_time < heartbeat->timestamp
                    && soft_state_[key]->local_time < local_time)
            {
                soft_state_[key]->local_time = local_time;
                soft_state_[key]->peer_time = heartbeat->timestamp;
            }
            else
            {
                // TODO
                std::cerr << "Fix me: "
                          << __LINE__ << " [ " << __FILE__  << "]"
                          << std::endl;
                exit(1);
            }
        }

    }

private:
    /// Read-write mutex for the soft state
    boost::shared_mutex& soft_state_mutex_;

    /// Soft state of the multicast membership
    kitty::soft_state& soft_state_;

    /// The underlying multicast receiver
    kitty::mcast_receiver mcast_receiver_;
};

} // namespace kitty

#endif // KITTY_MEMB_LISTENER_H

